from typing import Optional

from langgraph.graph import StateGraph, START, END
from pydantic import BaseModel


# from my_llm import llm


# 定义结构化状态
class MyState(BaseModel):
    x: int
    result: Optional[str] = None


def check_x(state: MyState) -> MyState:
    print(f"check_x, state: {state}")
    return state


# 判断x是奇数还是偶数，如果是奇数，则走handle_odd节点，否则走handle_even节点
# def is_even(state: MyState) -> bool:
#     return state.x % 2 == 0


def score_to_grade(state: MyState) -> str:
    grade: str = None
    if state.x >= 90:
        grade = "A"
    elif state.x >= 80:
        grade = "B"
    elif state.x >= 70:
        grade = "C"
    elif state.x >= 60:
        grade = "D"
    else:
        grade = "E"
    return grade


def handle_a(state: MyState) -> MyState:
    print(f"handle_a, x是90以上")
    return MyState(x=state.x, result="A")


def handle_b(state: MyState) -> MyState:
    print(f"handle_b, x是80以上")
    return MyState(x=state.x, result="B")


def handle_c(state: MyState) -> MyState:
    print(f"handle_c, x是70以上")
    return MyState(x=state.x, result="C")


def handle_d(state: MyState) -> MyState:
    print(f"handle_d, x是60以上")
    return MyState(x=state.x, result="D")


def handle_e(state: MyState) -> MyState:
    print(f"handle_e, x是60以下")
    return MyState(x=state.x, result="E")


# def handle_even(state: MyState) -> MyState:
#     print(f"handle_even, x是偶数")
#     return MyState(x=state.x, result="even")
#
#
# def handle_odd(state: MyState) -> MyState:
#     print(f"handle_odd, x是奇数")
#     return MyState(x=state.x, result="odd")


builder = StateGraph(MyState)
builder.add_node("check_x", check_x)
builder.add_node("handle_a", handle_a)
builder.add_node("handle_b", handle_b)
builder.add_node("handle_c", handle_c)
builder.add_node("handle_d", handle_d)
builder.add_node("handle_e", handle_e)

# 添加条件分支
builder.add_conditional_edges("check_x", score_to_grade, {
    "A": "handle_a",
    "B": "handle_b",
    "C": "handle_c",
    "D": "handle_d",
    "E": "handle_e",
})

# 衔接起始和结束
builder.add_edge(START, "check_x")
# builder.add_edge("handle_even", END)
# builder.add_edge("handle_odd", END)
builder.add_edge("handle_a", END)
builder.add_edge("handle_b", END)
builder.add_edge("handle_c", END)
builder.add_edge("handle_d", END)
builder.add_edge("handle_e", END)

graph = builder.compile()
